
# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from math import *
#from numpy import *
from client import n #the numpy package
#import numpy

def resub(D, rhs):
	
	""" solves D D^T = rhs by resubstituion.
		D is lower triangle-matrix from cholesky-decomposition """
	
	M = D.shape[0]
	x1= n.zeros((M,),float)
	x2= n.zeros((M,),float)

	# resub step 1
	for l in range(M): 
		sum = rhs[l]
		for j in range(l):
			sum -= D[l,j]*x1[j]
		x1[l] = sum/D[l,l]

	# resub step 2
	for l in range(M-1,-1,-1): 
		sum = x1[l]
		for j in range(l+1,M):
			sum -= D[j,l]*x2[j]
		x2[l] = sum/D[l,l]
	return x2

def calc_coeff(num_points, pol_degree, diff_order=0):
	
	""" calculates filter coefficients for symmetric savitzky-golay filter.
		see: http://www.nrbook.com/a/bookcpdf/c14-8.pdf
	
		num_points   means that 2*num_points+1 values contribute to the
			smoother.
	
		pol_degree   is degree of fitting polynomial
	
		diff_order   is degree of implicit differentiation.
			0 means that filter results in smoothing of function
			1 means that filter results in smoothing the first 
							derivative of function.
			and so on ...
	
	"""

	# setup normal matrix
	A = n.zeros((2*num_points+1, pol_degree+1), float)
	for i in range(2*num_points+1):
		for j in range(pol_degree+1):
			A[i,j] = pow(i-num_points, j)
	# calculate diff_order-th row of inv(A^T A)
	ATA = n.dot(A.transpose(), A)
	rhs = n.zeros((pol_degree+1,), float)
	rhs[diff_order] = 1
	D = n.linalg.cholesky(ATA)
	wvec = resub(D, rhs)

    # calculate filter-coefficients
	coeff = n.zeros((2*num_points+1,), float)
	for j in range(-num_points, num_points+1):
		x = 0.0
		for m in range(pol_degree+1):
			x += wvec[m]*pow(j, m)
		coeff[j+num_points] = x
	return coeff

def smooth(signal, coeff):
	
	""" applies coefficients calculated by calc_coeff()
	to signal """
	
	N = n.size(coeff-1)/2
	res = n.convolve(signal, coeff)
	return res[N:-N]

def savitzky_golay(data, kernel = 11, order = 4, deriv = 0):
    """
        applies a Savitzky-Golay filter
        input parameters:
        - data => data as a 1D numpy array
        - kernel => a positiv integer > 2*order giving the kernel size
        - order => order of the polynomal
        returns smoothed data as a numpy array

        invoke like:
        smoothed = savitzky_golay(<rough>, [kernel = value], [order = value]
    """
    #print kernel
    try:
            kernel = abs(int(kernel))
            order = abs(int(order))
    except ValueError, msg:
        raise ValueError("kernel and order have to be of type int (floats will be converted).")
    if kernel % 2 != 1 or kernel < 1:
        raise TypeError("kernel size must be a positive odd number, was: %d" % kernel)
    if kernel < order + 2:
        raise TypeError("kernel is to small for the polynomals\nshould be > order + 2")

    # a second order polynomal has 3 coefficients
    order_range = range(order+1)
    half_window = (kernel -1) // 2
    b = n.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)])
    # since we don't want the derivative, else choose [1] or [2], respectively
    m = n.linalg.pinv(b).A[deriv]
    window_size = len(m)
    half_window = (window_size-1) // 2

    # precompute the offset values for better performance
    offsets = range(-half_window, half_window+1)
    offset_data = zip(offsets, m)

    smooth_data = list()

    # temporary data, with padded zeros (since we want the same length after smoothing)
    data = n.concatenate((n.zeros(half_window), data, n.zeros(half_window)))
    for i in range(half_window, len(data) - half_window):
            value = 0.0
            for offset, weight in offset_data:
                value += weight * data[i + offset]
            smooth_data.append(value)
    return n.array(smooth_data)
